import numpy as np
import torch
import cv2
from tensorboardX import SummaryWriter
from torch.utils.data import DataLoader
# from dataset.data import DatasetPairCrop, restore, DatasetParamsMatting, DatasetMatting, DatasetPPM
from dataset.data import DatasetPPM


tb_logger = SummaryWriter('./logs/tb_img_pair')
# dataset = DatasetPairCrop(root='/home/chengk/params-photo-serving/chk/gaze_human_reason/raw/lock/')
# dataset = DatasetParamsMatting(root_img='/home/chengk/params-photo-serving/chk/gaze_human_reason/raw/lock/',
#                                 root_anno='/home/chengk/params-photo-serving/chk/gaze_human_reason/raw_matting')
# dataset = DatasetMatting(root='/home/chengk/chk/data/removebg-align/normal3',
#                         root_anno='/home/chengk/chk/data/removebg-align/normal3_removebg')
dataset = DatasetPPM(root='/home/chengk/chk/data/PPM-100/image/',
                    root_anno='/home/chengk/chk/data/PPM-100/matte/')
assert len(dataset) > 0, len(dataset)
dataloader = DataLoader(dataset, 1, num_workers=1, drop_last=False, shuffle=False)     # NOTE: please finish this function

def transform_inv(img_tensor):
    return (img_tensor.detach().cpu().numpy().transpose(1, 2, 0) + 1)/2

step = 0
for idx, data in enumerate(dataloader):
    '''

    img_anchor, bnd_box_anchor, img_target, bnd_box_target, size_ori = data
    # print(img_anchor.size(), img_target.size())
    # print(type(bnd_box_anchor), type(size_ori))
    # print(size_ori[0])
    # print(bnd_box_anchor)
    img_anchor = transform_inv(img_anchor[0])[..., 0]
    img_target = transform_inv(img_target[0])[..., 0]
    # print('transform_inv:', img_anchor.shape, img_target.shape)
    alpha_anchor, mask_anchor = restore(img_anchor, size_ori, bnd_box_anchor)
    alpha_target, mask_target = restore(img_target, size_ori, bnd_box_target)
    mask_overlap = (mask_anchor > 0) & (mask_target > 0)
    # print('restore:', alpha_anchor.shape, mask_anchor.shape)
    img_plot = np.concatenate([alpha_anchor[None, ...], mask_anchor[None, ...], 
                                alpha_target[None, ...], mask_target[None, ...],
                                mask_overlap[None, ...]], 0)
    img_plot = img_plot[:, None, ...]
    '''

    '''
    # DatasetPairCrop
    img_anchor, overlap_anchor, img_target, overlap_target = data
    img_anchor = transform_inv(img_anchor[0])[..., 0] # numpy HWC
    img_target = transform_inv(img_target[0])[..., 0] # numpy HWC
    if overlap_anchor is not None:
        h_anchor, w_anchor = img_anchor.shape[:2]
        h_target, w_target = img_target.shape[:2]

        x_min_anchor = int(w_anchor * overlap_anchor[0])
        x_max_anchor = int(w_anchor * overlap_anchor[1])
        y_min_anchor = int(h_anchor * overlap_anchor[2])
        y_max_anchor = int(h_anchor * overlap_anchor[3])
        
        x_min_target = int(w_target * overlap_target[0])
        x_max_target = int(w_target * overlap_target[1])
        y_min_target = int(h_target * overlap_target[2])
        y_max_target = int(h_target * overlap_target[3])

        # crop anchor & target region
        anchor_crop = img_anchor[y_min_anchor: y_max_anchor, x_min_anchor: x_max_anchor, ...]
        target_crop = img_target[y_min_target: y_max_target, x_min_target: x_max_target, ...]
        print('anchor crop:', anchor_crop.shape, '\t target crop:', target_crop.shape)
        # 打印两个ratio如果值相同表示宽高比是一致的
        ratio_anchor_crop = (y_max_anchor-y_min_anchor)/(x_max_anchor-x_min_anchor)
        ratio_target_crop = (y_max_target-y_min_target) / (x_max_target-x_min_target)
        print('anchor ratio:', ratio_anchor_crop, '\ttarget ratio:', ratio_target_crop)

        shape_resize = (256, int(256*(min(ratio_anchor_crop, ratio_target_crop))))
        tb_logger.add_images('image crop', np.concatenate([cv2.resize(anchor_crop, shape_resize)[None, ...],
                cv2.resize(target_crop, shape_resize)[None, ...]], 0)[:, None], step, dataformats='NCHW')
    '''
    '''
    img, alpha, receptive = data
    img = transform_inv(img[0])[..., 0] # input numpy HWC
    alpha = alpha.detach().numpy()
    receptive = receptive.detach().numpy()
    # print('img, alpha, receptive:', img.shape, alpha.shape, receptive.shape)
    '''
    img, alpha = data
    alpha = alpha.detach().numpy()
    print('alpha:', alpha.shape, alpha.dtype, '\timg:', img.shape)
    assert alpha.shape[-2:] == img.shape[-2:]

    img_plot = np.concatenate([(img+1)/2, 
                    np.concatenate([alpha]*3, 1),
                    ], 0)
    tb_logger.add_images('image pairs', img_plot, step, dataformats='NCHW')
    step += 1